#include "kernel_operator.h"

using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelXlogy
{
public:
    __aicore__ inline KernelXlogy() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t ALIGN_NUM, uint32_t core_size, 
                                uint32_t core_remain, uint32_t ub_data_num)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t blockIdx = GetBlockIdx();

        this->coreDataNum = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
        this->tileDataNum = ub_data_num;

        auto startPointer = core_size * GetBlockIdx();
        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1 + startPointer, this->coreDataNum);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x2 + startPointer, this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y + startPointer, this->coreDataNum);
        
        this->tileNum = this->coreDataNum / this->tileDataNum + (this->coreDataNum % this->tileDataNum > 0);


        pipe.InitBuffer(x1Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(x2Queue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
        pipe.InitBuffer(yQueue, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));
        pipe.InitBuffer(All0, this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(Bits, this->tileDataNum * sizeof(uint8_t));
    }

    __aicore__ inline void Process()
    {
        uint32_t loopCount = this->tileNum;
        for (int32_t i = 0; i < loopCount - 1; i++) 
        {
            CopyIn(i, this->tileDataNum);
            Compute(i, this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }

        uint32_t length = this->coreDataNum - this->tileDataNum * (loopCount - 1);
        CopyIn(loopCount - 1, length);
        Compute(loopCount - 1, length);
        CopyOut(loopCount - 1, length);
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) 
    {
        LocalTensor<DTYPE_X1> x1Local = x1Queue.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X1> x2Local = x2Queue.AllocTensor<DTYPE_X2>();

        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], length);
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], length);
        

        x1Queue.EnQue(x1Local);
        x2Queue.EnQue(x2Local);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_X1> x1Local = x1Queue.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = x2Queue.DeQue<DTYPE_X2>();

        LocalTensor<DTYPE_Y> yLocal = yQueue.AllocTensor<DTYPE_Y>();

        LocalTensor<uint8_t> bits = Bits.Get<uint8_t>();
        LocalTensor<DTYPE_X1> all0 = All0.Get<DTYPE_X1>();

        Compare(bits, x1Local, all0, CMPMODE::NE, length);
        Select(x2Local, bits, x2Local, (DTYPE_X1)1.0, SELMODE::VSEL_TENSOR_SCALAR_MODE, length);

        Ln(x2Local, x2Local, length);
        Mul(yLocal, x1Local, x2Local, length);

        yQueue.EnQue(yLocal);
        x1Queue.FreeTensor(x1Local);
        x2Queue.FreeTensor(x2Local);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<DTYPE_Y> yLocal = yQueue.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, length);
        yQueue.FreeTensor(yLocal);
    }

    GlobalTensor<DTYPE_X1> x1Gm, x2Gm;
    GlobalTensor<DTYPE_Y> yGm;

    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> x1Queue, x2Queue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;
    TBuf<QuePosition::VECCALC> Bits, All0;

    uint32_t coreDataNum = 0;
    uint32_t tileDataNum = 0;
    uint32_t tileNum = 0;
};

extern "C" __global__ __aicore__ void xlogy(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelXlogy op;
    op.Init(x1, x2, y, 
            tiling_data.ALIGN_NUM, tiling_data.core_size, 
            tiling_data.core_remain, tiling_data.ub_data_num);
    op.Process();
}